# coding=UTF-8
__author__ = 'Jakub Narloch'

from PyQt4.QtCore import *
from PyQt4.QtGui import *
from WEDT.Classifiers import NltkNaiveBayesClassifier, NltkDecisionTreeClassifier, DocumentClassifier, RandomForestDocumentClassifier
from WEDT.Core import PorterStemmerWordProcessor, LancasterStemmerWordProcessor, ProgressListener

class WorkThread(QThread):
    __documentClassifier = None

    def __init__(self, documentClassifier, parent=None):
        super(self.__class__, self).__init__(parent)
        self.__documentClassifier = documentClassifier

    def run(self):
        result = self.__documentClassifier.classify()

        self.emit(SIGNAL('workCompleted'), result)

    def workCompleted(self):
        pass

class ProgressDialog(QProgressDialog, ProgressListener):
    def __init__(self, label, min, max, parent=None):
        super(self.__class__, self).__init__(label, QString(), min, max, parent)
        self.setMinimumDuration(0)
        self.setWindowTitle(label)
        self.setModal(True)

        self.connect(self, SIGNAL('progressStarted'), self.handleProgressStarted)
        self.connect(self, SIGNAL('progress'), self.handleProgress)
        self.connect(self, SIGNAL('progressCompleted'), self.handleProgressCompleted)

    def progressStarted(self, label, max):
        pass

    def progress(self, current):
        pass

    def progressCompleted(self):
        pass

    def handleProgressStarted(self, label, max):
        self.show()
        self.setLabelText(label)
        self.setValue(0)
        self.setMaximum(max)

    def handleProgress(self, current):
        self.setValue(current)

    def handleProgressCompleted(self):
        self.hide()

    def notifyProgressStarted(self, label, max):
        self.emit(SIGNAL('progressStarted'), label, max)

    def notifyProgress(self, current):
        self.emit(SIGNAL('progress'), current)

    def notifyProgressCompleted(self):
        self.emit(SIGNAL('progressCompleted'))

class ResultDialog(QDialog):
    def __init__(self, result, parent = None):
        super(self.__class__, self).__init__(parent)
        self.setWindowTitle(u'Wyniki klasyfikacji')

        text = ""
        for documentResult in result.get_document_classes():
            text += ('%s: %s \n' % (documentResult[0], documentResult[1]))

        accuracyLabel = QLabel()
        accuracyLabel.setText(u'Dokładność klasyfikacji: %s' % result.get_accuracy())

        resultEdit = QTextEdit()
        resultEdit.setReadOnly(True)
        resultEdit.setText(text)

        layout = QVBoxLayout()
        layout.addWidget(accuracyLabel)
        layout.addWidget(resultEdit)
        self.setLayout(layout)

class MainWindow(QMainWindow):
    def __init__(self, parent=None):
        super(self.__class__, self).__init__(parent)
        self.setWindowTitle(u'Klasyfikator dokumentów')
        self.resize(400, 200)

        trainLabel = QLabel(u'Zbiór trenujący:')
        testLabel = QLabel(u'Zbiór testowy:')
        stemmerLabel = QLabel(u'Stemmer:')
        clarificationLabel = QLabel(u'Algorytm:')

        self.trainTextField = QLineEdit()
        self.testTextField = QLineEdit()

        self.stemmerComboBox = QComboBox()
        self.stemmerComboBox.addItems(['Porter', 'Lancaster'])

        self.algorithmComboBox = QComboBox()
        self.algorithmComboBox.addItems(['DecisionTree', 'NaiveBayes', 'RandomForest'])

        okButton = QPushButton(u'Klasyfikuj')
        browseTrainSetButton = QPushButton('...')
        browseTestSetButton = QPushButton('...')

        buttonLayout = QHBoxLayout()
        buttonLayout.addStretch()
        buttonLayout.addWidget(okButton)

        layout = QGridLayout()
        layout.addWidget(trainLabel, 0, 0)
        layout.addLayout(self.createQHBoxLayout(self.trainTextField, browseTrainSetButton), 0, 1)
        layout.addWidget(testLabel, 1, 0)
        layout.addLayout(self.createQHBoxLayout(self.testTextField, browseTestSetButton), 1, 1)
        layout.addWidget(stemmerLabel, 2, 0)
        layout.addWidget(self.stemmerComboBox, 2, 1)
        layout.addWidget(clarificationLabel, 3, 0)
        layout.addWidget(self.algorithmComboBox, 3, 1)
        layout.addLayout(buttonLayout, 4, 0, 1, 2)

        widget = QWidget()
        widget.setLayout(layout)
        self.setCentralWidget(widget)

        self.connect(okButton, SIGNAL('clicked()'), self.classify_documents)
        self.connect(browseTrainSetButton, SIGNAL('clicked()'), self.browse_train_set)
        self.connect(browseTestSetButton, SIGNAL('clicked()'), self.browse_test_set)

    def directory_chooser(self):
        dir = QFileDialog.getExistingDirectory(self, 'Folder',
                                               '.', QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks)

        return dir

    def createQHBoxLayout(self, *arg):
        layout = QHBoxLayout()

        for widget in arg:
            layout.addWidget(widget)

        return layout

    def browse_train_set(self):
        self.trainTextField.setText(self.directory_chooser())

    def browse_test_set(self):
        self.testTextField.setText(self.directory_chooser())

    def classify_documents(self):
        wordProcessor = None
        if self.stemmerComboBox.currentText() == 'Porter':
            wordProcessor = PorterStemmerWordProcessor()
        elif self.stemmerComboBox.currentText() == 'Lancaster':
            wordProcessor = LancasterStemmerWordProcessor()

        progressDialog = ProgressDialog('Klasyfikacja dokumentow.', 0, 100, self)
        progressDialog.setAutoReset(False)
        progressDialog.setAutoClose(False)

        if self.algorithmComboBox.currentText() == 'RandomForest':
            documentClassifier = RandomForestDocumentClassifier()
            documentClassifier.setWordProcessor(wordProcessor)
            documentClassifier.setTrainingSet(str(self.trainTextField.text()))
            documentClassifier.setTestSet(str(self.testTextField.text()))
            documentClassifier.setProgressListener(progressDialog)

            workThread = WorkThread(documentClassifier, self)
            self.connect(workThread, SIGNAL('workCompleted'), self.handleWorkCompleted)
            workThread.start()
        else:
            algorithm = None
            if self.algorithmComboBox.currentText() == 'DecisionTree':
                algorithm = NltkDecisionTreeClassifier()
            elif self.algorithmComboBox.currentText() == 'NaiveBayes':
                algorithm = NltkNaiveBayesClassifier()

            documentClassifier = DocumentClassifier()
            documentClassifier.setWordProcessor(wordProcessor)
            documentClassifier.setClassificationAlgorithm(algorithm)
            documentClassifier.setTrainingSet(str(self.trainTextField.text()))
            documentClassifier.setTestSet(str(self.testTextField.text()))
            documentClassifier.setProgressListener(progressDialog)

            workThread = WorkThread(documentClassifier, self)
            self.connect(workThread, SIGNAL('workCompleted'), self.handleWorkCompleted)
            workThread.start()

    def handleWorkCompleted(self, result):
        resultDialog = ResultDialog(result, self)
        resultDialog.show()